from scapy.all import *

class Comm(object):
    def __init__(self):
        self.iplist,self.maclist = readarptable()
        self.defaultfileds = ["dst_ip","ether_type"]
        self.threshold = 10

    def inject(**kwargs):
        pkt = gen_pkt(kwargs)
        sign = 0
        if pkt == None:
            print "Error,null pkt"
        else:
            rtt = []
            for i in range(1,3):
                ans, unans = srp(pkt,verbose=0,timeout=10)
                if(len(ans) != 0):
                    rx = ans[0][1]
                    tx = ans[0][0]
                    #print tx.sent_time
                    #print rx.time
                    rt = (rx.time - tx.sent_time)*1000
                    rtt.append(round(rt))
                else:
                    sign  = 33
                    #dropmem.append(m)
                    break
            if i != 1:
                delta = rtt[0] - rtt[-1]
                if delta > threshold:
                    sign = 22
                else:
                    sign = 11
        return sign
                #replymem[m] = rtt[0]
                #deltamem[m] = rtt[0] - rtt[-1]


    def gen_pkt(**kwargs):
        #check kwargs
        for i in self.defaultfileds:
            if kwargs.has_key(i) == false :
                print "Error! Can not generate pkt without enough information!"
                return None
        
        etype = kwargs.get("ether_type")
        dip = kwargs.get("dst_ip")
        dmac = 0
        #check dst_ip
        if dip not in self.iplist:
            print "not support to send pkt to an unkown host,check the ip bound"
            return None
        else:
            #check dst _mac
            dmac = maclist[iplist.index(dip)]
            if kwargs.has_key("dst_mac"):
                if kwargs.get("dst_mac") != dmac:
                    print "dst_ip is not at this dst_mac address,please check the args"
                    return None
        if etype == 0x0806:#arp
            pkt = Ether(dst = dmac,type = 0x0806)/ARP(pdst = dip,hwdst = dmac)
            return pkt
        else if etype == 0x0800:#ip
            eth = Ether(dst = dmac,type = 0x0800)
            tos = 0
            if kwargs.has_key("tos"):
                tos = kwargs.get("tos")
            if kwargs.has_key("proto"):
                proto = kwargs.get("proto")
                if proto == 1:#icmp
                    i_type = kwargs.get("icmp_type")
                    pkt = eth/IP(dst = dip,proto = 1, tos = tos)/ICMP(type=i_type)
                    return pkt
                else if proto == 6:
                    sp = kwargs.get("src_port")
                    dp = kwargs.get("dst_port")
                    pkt = eth/IP(dst= dip,proto = 6, tos = tos)/TCP(sport = sp, dport = dp)
                    return pkt
                else if proto == 17:
                    sp = kwargs.get("src_port")
                    dp = kwargs.get("dst_port")
                    pkt = eth/IP(dst= dip,proto = 6, tos = tos)/UDP(sport = sp, dport = dp)
                    return pkt 
         return pkt                        

 def readarptable(self):
    iplist = []
    maclist = []
    output = subprocess.Popen(["arp", "-n"],stdout=subprocess.PIPE)
        while True:
            line=output.stdout.readline()
            #print line
            if line !='':
                line_s = line.split(" ")
                #print line_s
                line_s = filter(lambda a: a !="",line_s)
                #print(line_s)

                if line_s[0][:1].isdigit() and line_s[2][:1].isdigit():
                    iplist.append(line_s[0])
                    maclist.append(line_s[2])
        else:
            break
    #print iplist
    #print maclist
    return iplist,maclist
    